import yaml
import matplotlib
import tkinter as tk
import pandas as pd
from tkinter import Text, Label, Button, END
from tkinter.ttk import Notebook, Combobox
from matplotlib.figure import Figure
from matplotlib.ticker import AutoLocator
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk
from matplotlib.backend_bases import key_press_handler
from tools.utils import loc_collection, loc_zq_collection
from cal_all import to_seasonal

matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 中文显示
matplotlib.rcParams['axes.unicode_minus'] = False  # 负号显示


def to_df(find_dic):
    data = pd.DataFrame(find_dic['value'], index=find_dic['t'], columns=find_dic['columns'])
    return data


def load_config():
    with open("config.yaml", 'r', encoding='utf-8') as f:
        data = yaml.load(f.read(), Loader=yaml.FullLoader)
    return data


def filter_cache_data(cache_data, name_):
    for data in cache_data:
        if name_ in data['columns']:
            return [data]
    return []


class Gui:
    def __init__(self, data_base, collection_name):
        super(Gui, self).__init__()
        self.root = tk.Tk()
        self.config = load_config()
        self.root.title("欢迎使用腾景人工智能大宗商品数据库")
        self.root.geometry("1860x960")
        self.data_base = data_base
        self.collection_name = collection_name
        self.bg = "#F07C82"
        self.seasonal = False
        self.three_dic = {}
        self.s_f = ('Microsoft YaHei', 22,)
        self.tree = self.ini_tree()
        self.book = Notebook(self.root, height=850, width=1508)
        self.book.place(x=310, y=0)
        self.fm_args = {"relief": "ridge", "bd": 3, "bg": self.bg}
        self.txt_args = {"bg": "#5F9EA0", 'font': self.s_f}
        self.fm_dic, self.img_dic = self.ini_fm()
        self.freq_combo, self.items_combo, self.year_combo = self.ini_combo()
        self.ini_label()
        self.button = self.ini_button()
        self.color_li = ['#CC0033', '#003399', '#FF6600', '#663366', '#7E2F8E', '#77AC30']

    def _ini_fig(self, fm):
        fig = Figure(figsize=(18, 10), dpi=78)
        canvas = FigureCanvasTkAgg(fig, fm)
        canvas.get_tk_widget().pack(side=tk.TOP, fill=tk.BOTH, expand=1)
        # spec = gridspec.GridSpec(ncols=3, nrows=6, figure=fig)
        toolbar = NavigationToolbar2Tk(canvas, fm)
        toolbar.update()
        canvas.mpl_connect("key_press_event", self._on_key_press1)
        return {'fig': fig, 'canvas': canvas, 'toolbar': toolbar}

    def _on_key_press1(self, event):
        print("you pressed {}".format(event.key))
        key_press_handler(event, self.fm_dic['show'], self.toolbar1)

    def ini_fm(self):
        fm_dic, fm_img = {}, {}
        for txt in ['show']:
            fm_dic[txt] = tk.Frame(self.book, height=780, width=1030, **self.fm_args)
            fm_dic[txt].place(x=10, y=10)
            fm_img[txt] = self._ini_fig(fm_dic[txt])
            self.book.add(fm_dic[txt], text=f" 【{txt}】 ")
        return fm_dic, fm_img

    def ini_combo(self):
        combo_freq = Combobox(self.fm_dic['show'], width=5)
        combo_freq['values'] = ['日', '周', '月']
        combo_freq.current(0)
        combo_freq.place(x=380, y=40)
        combo_freq.bind("<<ComboboxSelected>>", self.plot_fig)
        combo_items = Combobox(self.fm_dic['show'], width=45)
        combo_items['values'] = ['-', '-']
        combo_items.current(0)
        combo_items.place(x=10, y=40)
        combo_items.bind("<<ComboboxSelected>>", self.plot_fig)
        year_number_combo = Combobox(self.fm_dic['show'], width=5)
        year_number_combo['values'] = [str(i) for i in range(7)]
        year_number_combo.current(3)
        year_number_combo.place(x=540, y=40)
        year_number_combo.bind("<<ComboboxSelected>>", self.plot_fig)
        return combo_freq, combo_items, year_number_combo

    def ini_label(self):
        lb1 = Label(self.fm_dic['show'], text='请选择指标：', font=('Microsoft YaHei', 12,), bg='white')
        lb1.place(x=10, y=10)
        lb2 = Label(self.fm_dic['show'], text='频率:', font=('Microsoft YaHei', 12,), bg='white')
        lb2.place(x=380, y=10)
        lb3 = Label(self.fm_dic['show'], text='年份', font=('Microsoft YaHei', 12,), bg='white')
        lb3.place(x=540, y=10)

    def ini_button(self):
        button = Button(self.fm_dic['show'], width=10, text="季节性切换", command=self.seasonal_change)
        button.place(x=450, y=30)
        return button

    def ini_tree(self):
        tree = tk.ttk.Treeview(self.root, height=41, selectmode="extended")
        tree.column("#0", minwidth=0, width=300, )
        tree.place(x=0, y=25)
        tree_dict = {}
        coll_ = self._loc_col()
        for first_i, (tjd_category, chinese_names) in enumerate(self.config['tjd_category'].items()):
            tree_dict[tjd_category] = tree.insert("", first_i, tjd_category, text=tjd_category, values='1')
            for second_i, chinese_name in enumerate(chinese_names):
                f1_str = f"{tjd_category}>>{chinese_name}"
                tree_dict[f1_str] = tree.insert(tree_dict[tjd_category], second_i, f1_str, text=chinese_name,
                                                values='2')
                for third_i, mode in enumerate(self.config['pages']):
                    f2_str = f"{tjd_category}>>{chinese_name}>>{mode}"
                    tree_dict[f2_str] = tree.insert(tree_dict[f1_str], third_i, f2_str, text=mode, values='3')
                    find_dic = {'chinese_name': chinese_name, 'mode': mode, 'seasonal': False}
                    fit_data = coll_.find(find_dic)
                    item_dic = self.three_dic.setdefault(f2_str, {})
                    print(f2_str)
                    for fouth_i, d in enumerate(fit_data):
                        freq = d['freq']
                        name_dic = item_dic.setdefault(d['columns'][0], {})
                        freq_dic = name_dic.setdefault(freq, {})
                        freq_dic[freq] = d
        tree.bind("<<TreeviewSelect>>", self._tree_fun)
        return tree

    def _loc_col(self):
        if self.data_base == 'zq':
            coll = loc_zq_collection('to_it', self.collection_name)
        else:
            coll = loc_collection('to_it', self.collection_name)
        return coll

    def _tree_fun(self, event):
        name = event.widget.selection()[0]
        self.name = name
        if name in self.three_dic:
            self.tjd_category, self.chinese_name, self.mode = self.name.split('>>')
            find_dic = {'chinese_name': self.chinese_name, 'mode': self.mode, 'pred': False}
            v = self._loc_col().find(find_dic)
            names = [d['columns'][0] for d in v]
            names = list(set(names))
            print('names', names)
            self.items_combo['value'] = names
            self.items_combo.set(names[0])

    def plot_fig(self, event):
        freq = self.freq_combo.get()
        name_ = self.items_combo.get()
        dic_ = {'chinese_name': self.chinese_name, 'mode': self.mode, 'freq': freq, 'pred': False}
        dic_pred = {'chinese_name': self.chinese_name, 'mode': self.mode, 'freq': freq, 'pred': True}
        print('looking data in mongo', dic_)
        cache_data = list(self._loc_col().find(dic_))
        cache_data_pred = list(self._loc_col().find(dic_pred))
        print('fited data', len(cache_data), 'pred_fited_data', len(cache_data_pred))
        if len(cache_data) > 0:
            fit_data = filter_cache_data(cache_data, name_)
            fit_data_pred = filter_cache_data(cache_data_pred, name_)
            fig = self.img_dic['show']['fig']
            fig.clear()
            ax = fig.add_subplot(111)
            df = to_df(fit_data[0])
            if len(fit_data_pred) > 0:
                pred = True
                df_pred = to_df(fit_data_pred[0])
            else:
                pred = False
            if self.seasonal:
                df = to_seasonal(df, self.mode, freq=freq)
                year_num = int(self.year_combo.get())
                df = df.iloc[:, -min(df.shape[1], year_num):]
                if pred:
                    df_pred = to_seasonal(df_pred, self.mode, freq=freq)
            else:
                df = df.iloc[:, [0]]
                df.index = pd.to_datetime(df.index)
                if pred:
                    df_pred.index = pd.to_datetime(df_pred.index)
            for i, (name, se) in enumerate(df.iteritems()):
                if i + 1 < df.shape[1]:
                    ax.plot(se, label=name, color=self.color_li[i])
                else:
                    ax.plot(se, label=name, color='black', lw=2)
                if pred:
                    ax.plot(df_pred, color='green', lw=4, label='pred')
            ax.legend()
            ax.grid()
            ax.xaxis.set_major_locator(AutoLocator())
            self.img_dic['show']['canvas'].draw()

    def seasonal_change(self):
        if self.seasonal:
            self.seasonal = False
        else:
            self.seasonal = True
        self.plot_fig('-')


if __name__ == "__main__":
    g = Gui('wuwei', '2022-01-17')
    g.root.mainloop()
